"""
    Computational Economics
    2.8: Schelling's Segregation Model
    http://johnstachurski.net/lectures/schelling.html

    REFERENCES
    http://en.wikipedia.org/wiki/Thomas_Schelling
    http://jkwiens.com/2010/01/02/creating-movies-with-pyplotpylab/
    http://matplotlib.sourceforge.net/examples/animation/movie_demo.html
    http://mindseye.no/2011/04/03/animating-plots-and-waves-in-python-using-matplotlib/
    http://stackoverflow.com/questions/4092927/

    OUTLINE OF MODEL
    Two populations of agents: red and blue

    Initially, they are integrated (randomly distributed).

    Each round, each agent is given a chance to stay or move:
    - If happy, stay
    - If unhappy, move in the direction of a similar agent type

    Agent is happy when:
    - Half or more of 10 nearest neighbors are same type
    - Nearest is in terms of Euclidean distance

    Continuing until no-one wishes to move

    MODELING
    Agents are modeled as objects

        * Data:
              o type and location
        * Methods:
              o is_happy
              o move

    Rules for moving:
        * move 1 unit
        * move in direction of a neighbor

    Pseudocode for the main loop

        while agents are still moving:
            for agent in agents:
                give agent the opportunity to move

    Use 200 agents of each type

    Plot the map continuously to create an animated movie

"""
import matplotlib
#matplotlib.use('Agg')

import sys, os
from os.path import (exists, dirname, join)
import random
import pylab
import unittest
import math
import time
from datetime import datetime
import subprocess



#
# Globals
#
AGENT_SIZE = 6


class Agent:
    type = ''
    start_location = [0,0]
    location = [0, 0]
    locations = []
    mobility = 5
    color = 'white'
    hilite = 'yellow'
    size = AGENT_SIZE

    def __init__(self, type=None, location=None):
        if type:
            self.type = type

        if location:
            self.start_location = location

        self.location = self.start_location
        self.num_relocations = 0

    def __str__(self):
        return '<Agent id=%s type=%s location=%s relocations=%s>' % (
            id(self), self.type, tuple(self.location), len(self.locations) )

    def is_happy(self, neighbors, min_same_type):
        same_type = [1 for n, d in neighbors if n.type == self.type]
        return sum(same_type) >= min_same_type

    def move(self, neighbors, board_dim):
        return self.move_randomly(board_dim)

    def move_randomly(self, board_dim):
        # randomly move in x, y directions
        x, y = self.location
        dx, dy = (random.uniform(-self.mobility,self.mobility),
                  random.uniform(-self.mobility,self.mobility))
        xnew, ynew = x + dx, y + dy

        # stop at board boundaries
        xmin, xmax, ymin, ymax = board_dim
        xnew = max([xnew, xmin+self.size])
        xnew = min([xnew, xmax-self.size])
        ynew = max([ynew, ymin+self.size])
        ynew = min([ynew, ymax-self.size])

        self.location = (xnew,ynew)
        self.locations.append(self.location)

    def move_toward(self, neighbors, board_dim):
        # pick nearest neighbor
        like_neighbors = [n for n,d in neighbors if n.type == self.type]

        # move in direction of random life neighbord
        if like_neighbors:
            target = random.choice(like_neighbors)
        else:
            return

        # get new location
        new_location = self.move_from_point_toward_point(self.location,
            target.location, self.mobility)

        # stop at board boundaries
        xmin, xmax, ymin, ymax = board_dim
        x = max([new_location[0], xmin-self.size])
        x = min([new_location[0], xmax-self.size])
        y = max([new_location[1], ymin-self.size])
        y = min([new_location[1], ymax-self.size])

        self.location = (x,y)
        self.locations.append(self.location)

    def move_from_point_toward_point(self, p1, p2, d):
        x1, y1 = p1
        x2, y2 = p2

        dx = float(x2 - x1)
        dy = float(y2 - y1)

        # avoid overshoot
        distance = math.sqrt(math.pow(dx,2) + math.pow(dy,2))
        if d > distance:
            d = distance - 1

        if dx > 0:
            xd = 1
        elif dx < 0:
            xd = -1
        else:
            xd = 0

        if dy > 0:
            yd = 1
        elif dy < 0:
            yd = -1
        else:
            yd = 0

        if not xd:
            return (x1, y1 + (d * yd))

        if not yd:
            return (x1 + (d *xd), y1)

        theta = math.atan(abs(dy)/abs(dx))
        dx = d * math.cos(theta)
        dy = d * math.sin(theta)

        new_x = x1 + (dx * xd)
        new_y = y1 + (dy * yd)

        return new_x, new_y

    def measure_distance(self, location):
        x1, y1 = location
        x0, y0 = self.location
        return math.sqrt((x1-x0)**2 + (y1-y0)**2)

    def get_random_location(self):
        return (random.uniform(0, 1), random.uniform(0, 1))

class RedAgent(Agent):
    type = 'Red'
    color = '#a80000'
    normal = '#a80000'
    hilite = '#ff0000'

class BlueAgent(Agent):
    type = 'Blue'
    color = '#0000a8'
    normal = '#0000a8'
    hilite = '#0000ff'


class Simulation:

    # parameters
    num_red_agents = 50
    num_blue_agents = 50
    num_neighbors = 10
    min_same_neighbors = 5
    max_rounds = 0
    debug = False
    record = False

    # movie parameters
    movie_title = 'movie'
    title = 'Animated Schelling Simulation'
    image_dir = '/tmp/schelling'
    show = True
    record = True

    # internal parameters
    board_dims = [-100,100,-100,100]        # xmin, xmax, ymin, ymax
    round = 0
    moves_last_round = 1
    agents = []

    def __init__(self):
        pylab.ion()

        # prepare output dir
        if not exists(self.image_dir):
            os.mkdir(self.image_dir)

        # delete files in working dir
        self.delete_all_png_files()

        self.setup_game()

    def setup_game(self):
        # create teams
        for n in range(self.num_blue_agents):
            location = self.get_random_board_location()
            self.agents.append(BlueAgent(location=location))

        for n in range(self.num_red_agents):
            location = self.get_random_board_location()
            self.agents.append(RedAgent(location=location))

    def play_game(self):
        while not self.is_finished():
            self.round += 1
            t = time.time()
            self.moves_last_round = self.play_round()
            self.draw_map()
            tr = time.time() - t
            print '[%.2fs] ROUND %04d: %s of %s agents moved' % (tr,
                self.round, self.moves_last_round, len(self.agents))

        self.create_movie()

    def play_round(self):
        num_movers = 0

        random.shuffle(self.agents)
        for agent in self.agents:
            neighbors = self.find_agent_neighbors(agent, self.num_neighbors)
            is_happy = agent.is_happy(neighbors, self.min_same_neighbors)
            if is_happy:
                agent.color = agent.normal
            if not is_happy:
                agent.color = agent.hilite
                agent.move(neighbors, self.board_dims)
                num_movers += 1
            self.d('agent %s turn in round %s: %s' % (id(agent), self.round,
                is_happy))

        return num_movers

    def find_agent_neighbors(self, agent, num):
        neighbor_distances = []
        for n in self.agents:
            if n is agent:
                continue
            d = self.measure_distance(agent.location, n.location)
            neighbor_distances.append((n, d))
        closest_neighbors = sorted(neighbor_distances, key=lambda tup: tup[1], reverse=False)
        return closest_neighbors[:num]

    def measure_distance(self, coords1, coords2):
        """coords should be an x, y tuple"""
        x0, y0 = coords1
        x1, y1 = coords2
        return math.sqrt((x1-x0)**2 + (y1-y0)**2)

    def get_random_board_location(self):
        xmin, xmax, ymin, ymax = self.board_dims
        return (random.uniform(xmin, xmax), random.uniform(ymin, ymax))

    def draw_board(self):
        #self.pause('drawing grid', 3)
        pylab.grid(True)
        pylab.title(self.title)

        # v = [xmin, xmax, ymin, ymax]
        pylab.axis(self.board_dims)

    def draw_map(self):
        #self.pause('clearing board', 3)
        pylab.clf()
        X,Y = 'x','y'

        #self.pause('drawing board', 3)
        # set player locations
        groups = {}
        for agent in self.agents:
            x, y = agent.location
            if agent.color not in groups:
                groups[agent.color] = {X:[], Y:[]}

            groups[agent.color][X].append(x)
            groups[agent.color][Y].append(y)

        # draw board
        self.draw_board()

        # plot each group
        for color in groups:
            #print color, len(groups[color][X]), len(groups[color][Y])
            pylab.plot(groups[color][X], groups[color][Y], marker='o',
                markerfacecolor=color, markersize=AGENT_SIZE, linestyle='None')

        pylab.axis(self.board_dims)
        img_name = 'frame%06d.png' % (self.round)
        img_path = join(self.image_dir, img_name)

        if self.show:
            pylab.draw()
        #time.sleep(3)

        if self.record:
            pylab.savefig(img_path)

    def is_finished(self):
        if self.max_rounds and (self.round >= self.max_rounds):
            return True
        if not self.moves_last_round:
            return True
        return False

    def create_movie(self):
        """ffmpeg -r 10 -sameq -i /tmp/schelling/frame%06d.png /tmp/schelling/test.mp4"""
        output_file = join(self.image_dir, '%s.mp4' % (self.movie_title))

        command = (
            'ffmpeg',
            '-r', '10',
            '-sameq',
            '-y',
            '-i', join(self.image_dir, 'frame%06d.png'),
            output_file
        )

        print "\n\ncreating movie:\n%s\n\n" % ' '.join(command)
        subprocess.check_call(command)
        print "movie created:", output_file
        self.delete_all_png_files()

    def delete_all_png_files(self):
        for f in os.listdir(self.image_dir):
            file_path = os.path.join(self.image_dir, f)
            if file_path.endswith('.png'):
                os.unlink(file_path)

    def d(self, message):
        if self.debug:
            print '[s] %s' % (time.time(), message)

    def pause(self, message, seconds=1):
        print message
        time.sleep(seconds)



def main():
    game = Simulation()
    game.title = 'Schelling Segregation Model'
    game.max_rounds = 500
    game.show = True
    game.movie_title = '%s-r%s' % (datetime.strftime(datetime.now(), '%Y%m%d'), game.max_rounds)
    game.play_game()
    print 'game over at round', game.round

#
# MAIN
#
if __name__ == "__main__":
    main()
    print '%s: ok' % (__file__)
